import os
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
from .train import GarbageClassifier

class GarbagePredictor:
    def __init__(self, model_path, device='cuda' if torch.cuda.is_available() else 'cpu'):
        self.device = device
        self.model = GarbageClassifier(num_classes=6)
        self.model.load_state_dict(torch.load(model_path, map_location=device))
        self.model.to(device)
        self.model.eval()
        
        # 原始类别到目标类别的映射
        self.source_classes = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
        self.target_classes = ['recyclable', 'harmful', 'wet', 'dry']
        self.target_classes_zh = ['可回收物', '有害垃圾', '湿垃圾', '干垃圾']
        self.target_class_mapping = dict(zip(self.target_classes, self.target_classes_zh))
        
        # 定义类别映射关系
        self.class_mapping = {
            'cardboard': 'recyclable',
            'glass': 'recyclable',
            'metal': 'recyclable',
            'paper': 'recyclable',
            'plastic': 'recyclable',
            'trash': 'dry'
        }
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    
    def predict(self, image_path):
        # 加载并预处理图像
        image = Image.open(image_path).convert('RGB')
        image_tensor = self.transform(image).unsqueeze(0).to(self.device)
        
        # 进行预测
        with torch.no_grad():
            outputs = self.model(image_tensor)
            _, preds = torch.max(outputs, 1)
            source_class = self.source_classes[preds.item()]
            pred_class = self.class_mapping[source_class]
            pred_class_zh = self.target_class_mapping[pred_class]
            
            # 计算原始类别的概率
            probs = torch.nn.functional.softmax(outputs, dim=1)[0]
            source_probs = {cls: float(prob) for cls, prob in zip(self.source_classes, probs)}
            
            # 将原始类别的概率映射到目标类别
            target_probs = {cls: 0.0 for cls in self.target_classes_zh}
            for src_cls, prob in source_probs.items():
                target_cls = self.class_mapping[src_cls]
                target_cls_zh = self.target_class_mapping[target_cls]
                target_probs[target_cls_zh] += prob
        
        return {
            'predicted_class': pred_class_zh,
            'probabilities': target_probs,
            'image_data': image_path
        }

def main():
    # 示例用法
    model_path = 'best_model.pth'
    predictor = GarbagePredictor(model_path)
    
    # 测试图像预测
    test_image = 'path/to/test/image.jpg'
    if os.path.exists(test_image):
        result = predictor.predict(test_image)
        print(f"预测类别: {result['predicted_class']}")
        print("各类别概率:")
        for cls, prob in result['probabilities'].items():
            print(f"{cls}: {prob:.4f}")

if __name__ == '__main__':
    main()